import os
import argparse
 
def batch_rename(directory, prefix="image", padding=3, start=1, extensions=None, output_file="datasets.txt"):
    """
    批量重命名图片文件并生成文件名列表
    参数：
    directory: 目标目录路径
    prefix: 新文件名前缀（默认：image）
    padding: 序号位数（默认：3）
    start: 起始序号（默认：1）
    extensions: 要处理的文件扩展名列表（默认：['jpg', 'jpeg', 'png', 'gif', 'webp']）
    output_file: 输出文件名列表的文件路径（默认：datasets.txt）
    """
    if not os.path.isdir(directory):
        print(f"错误：目录 {directory} 不存在")
        return
 
    # 默认支持的图片格式
    default_exts = ['jpg', 'jpeg', 'png', 'gif', 'webp']
    extensions = extensions if extensions else default_exts
    extensions = [ext.lower() for ext in extensions]
 
    # 获取所有图片文件
    files = []
    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)
        if os.path.isfile(file_path):
            ext = filename.split('.')[-1].lower()
            if ext in extensions:
                files.append(file_path)
 
    if not files:
        print("没有找到可处理的图片文件")
        return
 
    # 按修改时间排序（可根据需要修改排序方式）
    files.sort(key=lambda x: os.path.getmtime(x))
 
    # 准备输出文件
    output_path = os.path.join(directory, output_file)
    
    # 开始重命名并生成列表
    counter = start
    with open(output_path, 'w') as f:
        for old_path in files:
            dir_name = os.path.dirname(old_path)
            filename = os.path.basename(old_path)
            ext = filename.split('.')[-1].lower()
            
            # 生成新文件名
            new_name = f"{prefix}_{counter:0{padding}d}.{ext}"
            new_path = os.path.join(dir_name, new_name)
 
            # 处理文件名冲突
            while os.path.exists(new_path):
                print(f"警告：文件 {new_name} 已存在，自动跳过")
                counter += 1
                new_name = f"{prefix}_{counter:0{padding}d}.{ext}"
                new_path = os.path.join(dir_name, new_name)
 
            # 执行重命名
            os.rename(old_path, new_path)
            print(f"重命名：{filename} -> {new_name}")
            
            # 写入文件名列表
            f.write(f"{new_name}\n")
            
            counter += 1
 
    print(f"成功处理 {len(files)} 个文件，文件名列表已保存至 {output_path}")
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="批量重命名图片文件并生成文件名列表")
    parser.add_argument("directory", help="目标目录路径")
    parser.add_argument("-p", "--prefix", default="image", help="文件名前缀（默认：image）")
    parser.add_argument("-n", "--padding", type=int, default=3, help="序号位数（默认：3）")
    parser.add_argument("-s", "--start", type=int, default=1, help="起始序号（默认：1）")
    parser.add_argument("-e", "--extensions", nargs='+', 
                      help="要处理的文件扩展名（默认：jpg, jpeg, png, gif, webp）")
    parser.add_argument("-o", "--output", default="datasets.txt", 
                      help="输出文件名列表的文件路径（默认：datasets.txt）")
    
    args = parser.parse_args()
    
    batch_rename(
        directory=args.directory,
        prefix=args.prefix,
        padding=args.padding,
        start=args.start,
        extensions=args.extensions,
        output_file=args.output
    )